from ltp import LTP
ltp = LTP()
sents = ltp.sent_split(["他叫汤姆去拿外衣。", "汤姆生病了。他去了医院。"])
print(sents)

# user_dict.txt 是词典文件， max_window是最大前向分词窗口
ltp.init_dict(path="user_dict.txt", max_window=4)
# 也可以在代码中添加自定义的词语
ltp.add_words(words=["负重前行", "长江大桥"], max_window=4)

segment, _ = ltp.seg(["他叫汤姆去拿外衣。"])
# [['他', '叫', '汤姆', '去', '拿', '外衣', '。']]

print(segment)
# 对于已经分词的数据
# segment, hidden = ltp.seg(["他/叫/汤姆/去/拿/外衣/。".split('/')])


seg, hidden = ltp.seg(["他叫汤姆去拿外衣。"])

print(seg)
# print(hidden)

pos = ltp.pos(hidden)
print(pos)
# [['他', '叫', '汤姆', '去', '拿', '外衣', '。']]
# [['r', 'v', 'nh', 'v', 'v', 'n', 'wp']]

ner = ltp.ner(hidden)
# [['他', '叫', '汤姆', '去', '拿', '外衣', '。']]
# [[('Nh', 2, 2)]]

tag, start, end = ner[0][0]
print(tag,":", "".join(seg[0][start:end + 1]))


srl = ltp.srl(hidden)
# [['他', '叫', '汤姆', '去', '拿', '外衣', '。']]
# [
#     [
#         [],                                                # 他
#         [('ARG0', 0, 0), ('ARG1', 2, 2), ('ARG2', 3, 5)],  # 叫 -> [ARG0: 他, ARG1: 汤姆, ARG2: 拿外衣]
#         [],                                                # 汤姆
#         [],                                                # 去
#         [('ARG0', 2, 2), ('ARG1', 5, 5)],                  # 拿 -> [ARG0: 汤姆, ARG1: 外衣]
#         [],                                                # 外衣
#         []                                                 # 。
#     ]
# ]
srl = ltp.srl(hidden, keep_empty=False)
# [
#     [
#         (1, [('ARG0', 0, 0), ('ARG1', 2, 2), ('ARG2', 3, 5)]), # 叫 -> [ARG0: 他, ARG1: 汤姆, ARG2: 拿外衣]
#         (4, [('ARG0', 2, 2), ('ARG1', 5, 5)])                  # 拿 -> [ARG0: 汤姆, ARG1: 外衣]
#     ]
# ]

dep = ltp.dep(hidden)
# [['他', '叫', '汤姆', '去', '拿', '外衣', '。']]
# [
#     [
#         (1, 2, 'SBV'),
#         (2, 0, 'HED'),    # 叫 --|HED|--> ROOT
#         (3, 2, 'DBL'),
#         (4, 2, 'VOB'),
#         (5, 4, 'COO'),
#         (6, 5, 'VOB'),
#         (7, 2, 'WP')
#     ]
# ]

sdp = ltp.sdp(hidden, graph=False)
# [['他', '叫', '汤姆', '去', '拿', '外衣', '。']]
# [
#     [
#         (1, 2, 'Agt'),
#         (2, 0, 'Root'),   # 叫 --|Root|--> ROOT
#         (3, 2, 'Datv'),
#         (4, 2, 'eEfft'),
#         (5, 4, 'eEfft'),
#         (6, 5, 'Pat'),
#         (7, 2, 'mPunc')
#     ]
# ]

sdp = ltp.sdp(hidden, graph=True)
print(sdp)
# [['他', '叫', '汤姆', '去', '拿', '外衣', '。']]
# [
#     [
#         (1, 2, 'Agt'),
#         (2, 0, 'Root'),   # 叫 --|Root|--> ROOT
#         (3, 2, 'Datv'),
#         (3, 4, 'Agt'),
#         (3, 5, 'Agt'),
#         (4, 2, 'eEfft'),
#         (5, 4, 'eEfft'),
#         (6, 5, 'Pat'),
#         (7, 2, 'mPunc')
#     ]
# ]